#!/usr/bin/env python
# encoding=utf-8


import os,pprint
import sys, codecs

work_dir = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../../'))
print work_dir
sys.path.insert(0, work_dir)
pprint.pprint(sys.path)

from matchzoo.inputs.preparation import Preparation
from matchzoo.inputs.preprocess import Preprocess, NgramUtil

srcdir = './'
dstdir = './'

stopwords_path = os.path.join('../', 'stopwords.txt')

single_word = True

infiles = [ srcdir + 'ads_train_mz.txt', srcdir + 'ads_valid_mz.txt', srcdir + 'ads_test_mz.txt']

def read_dict(infile):
    word_dict = {}
    for line in codecs.open(infile, 'r', encoding='utf-8'):
        r = line.strip().split()
        word_dict[r[1]] = r[0]
    return word_dict
def read_doc(infile):
    doc = {}
    for line in codecs.open(infile, 'r', encoding='utf-8'):
        r = line.strip().split()
        doc[r[0]] = r[1:]
        #assert len(doc[r[0]]) == int(r[1])
    return doc
def filter_triletter(tri_stats, min_filter_num=5, max_filter_num=10000):
    tri_dict = {}
    tri_stats = sorted(tri_stats.items(), key=lambda d:d[1], reverse=True)
    for triinfo in tri_stats:
        if triinfo[1] >= min_filter_num and triinfo[1] <= max_filter_num:
            if triinfo[0] not in tri_dict:
                tri_dict[triinfo[0]] = len(tri_dict)
    return tri_dict

if __name__ == '__main__':
    prepare = Preparation()

    corpus, rel_train, rel_valid, rel_test = prepare.run_with_train_valid_test_corpus(infiles[0], infiles[1], infiles[2])
    print('total corpus : %d ...' % (len(corpus)))
    print('total relation-train : %d ...' % (len(rel_train)))
    print('total relation-valid : %d ...' % (len(rel_valid)))
    print('total relation-test: %d ...' % (len(rel_test)))
    prepare.save_corpus(dstdir + 'corpus.txt', corpus)

    prepare.save_relation(dstdir + 'relation_train.txt', rel_train)
    prepare.save_relation(dstdir + 'relation_valid.txt', rel_valid)
    prepare.save_relation(dstdir + 'relation_test.txt', rel_test)
    print('Preparation finished ...')

    stopwords = []
    with codecs.open(stopwords_path, encoding='utf-8') as f_stopwords:
        for line in f_stopwords.readlines():
            stopwords.append(line.strip())
    preprocessor = Preprocess(word_stem_config={'enable': False}, word_filter_config={'min_freq': 2, 'stop_words': set(stopwords)})
    dids, docs = preprocessor.run(dstdir + 'corpus.txt')
    preprocessor.save_word_dict(dstdir + 'word_dict.txt', True)
    preprocessor.save_words_stats(dstdir + 'word_stats.txt', True)

    fout = codecs.open(dstdir + 'corpus_preprocessed.txt', 'w', encoding='utf-8')
    for inum, did in enumerate(dids):
        fout.write('%s %s %s\n' % (did, len(docs[inum]), ' '.join(map(str, docs[inum]))))
    fout.close()
    print('Preprocess finished ...')

    #dssm_corp_input = dstdir + 'corpus_preprocessed.txt'
    #dssm_corp_output = dstdir + 'corpus_preprocessed_dssm.txt'
    word_dict_input = dstdir + 'word_dict.txt'
    triletter_dict_output = dstdir + 'triletter_dict.txt'
    word_triletter_output = dstdir + 'word_triletter_map.txt'
    word_dict = read_dict(word_dict_input)
    word_triletter_map = {}
    triletter_stats = {}
    for wid, word in word_dict.items():
        if not single_word:
            nword = '#' + word + '#'
            ngrams = NgramUtil.ngrams(list(nword), 3, '')
        else:
            ngrams = [word]
        word_triletter_map[wid] = []
        for tric in ngrams:
            if tric not in triletter_stats:
                triletter_stats[tric] = 0
            triletter_stats[tric] += 1
            word_triletter_map[wid].append(tric)
    triletter_dict = filter_triletter(triletter_stats, 0, 10000)
    with codecs.open(triletter_dict_output, 'w', encoding='utf-8') as f:
        for tri_id, tric in triletter_dict.items():
            line = tri_id + ' ' + str(tric) + '\n'
            f.write(line)
    with codecs.open(word_triletter_output, 'w', encoding='utf-8') as f:
        for wid, trics in word_triletter_map.items():
            line = str(wid) + ' ' + ' '.join([str(triletter_dict[k]) for k in trics if k in triletter_dict]) + '\n'
            f.write(line)
    print('Triletter Processing finished ...')

